from funasr import AutoModel

def check_vad(vad_list):
    # 检查列表长度是否是2的整数倍
    if len(vad_list) % 2 == 0:
        # 获取最后两个子列表
        last_two = vad_list[-2:]
        # 分别取每个子列表的最大值
        max_values = [max(sublist) for sublist in last_two]
        # 计算两个最大值之间的绝对差值
        result = abs(max_values[1] - max_values[0])
        return result
    else:
        # 如果不是2的整数倍，可以抛出异常或返回特定值表示错误
        return 0

# 这个不用改，表示单个音频的长度.
chunk_size = 50 # ms
model = AutoModel(model="/root/yuehu/assets/speech_fsmn_vad_zh-cn-16k-common-pytorch", device='cuda:1')

import soundfile

# wav_file = f"{model.model_path}/example/vad_example.wav"
wav_file = 'new.wav'
speech, sample_rate = soundfile.read(wav_file)
chunk_stride = int(chunk_size * sample_rate / 1000)

cache = {}

vad_res = []
total_chunk_num = int(len((speech)-1)/chunk_stride+1)
for i in range(total_chunk_num):
    speech_chunk = speech[i*chunk_stride:(i+1)*chunk_stride]
    is_final = i == total_chunk_num - 1
    res = model.generate(input=speech_chunk, cache=cache, is_final=is_final, chunk_size=chunk_size, max_end_silence_time=400, max_start_silence_time=1000)#, output_dir='output_dir/'
    if len(res) == 0:
        continue
    value = res[0]['value']
    if len(value) == 0:
        continue
    vad_res.append(res[0]['value'][0])

    #表示至少识别 1 秒的长度才认为需要结束.
    if check_vad(vad_res) > 1000:
        break

print(vad_res)